import torch
import cv2 as cv
import time

model = torch.hub.load("./", "custom", path="m_Final_v23.pt", source='local')
model.conf = 0.7
model.iou = 0.4

det = cv.VideoCapture(0)
det.set(cv.CAP_PROP_FRAME_WIDTH, 1280)
det.set(cv.CAP_PROP_FRAME_HEIGHT, 960)
det.set(cv.CAP_PROP_EXPOSURE, 15000)
print(det.get(3), det.get(4), det.get(cv.CAP_PROP_FPS), det.get(cv.CAP_PROP_EXPOSURE))

while True:
    ret, frame = det.read()
    if not ret:
        print("Error: Unable to read frame.")
        det.release()
        break
    else:
        start = time.time()
        height, width, _ = frame.shape

        # # 计算裁剪的起始坐标
        x = (width - 640) // 2
        y = (height - 560) // 2

        # 进行中心裁剪
        frame = frame[y:y + 640, x:x + 640]
        frame = cv.cvtColor(frame, cv.COLOR_BGR2RGB)
        results = model(frame)
        # results = self.model(median)
        image = results.render()[0]
        image = cv.cvtColor(image, cv.COLOR_RGB2BGR)
        df = results.pandas().xyxy[0].sort_values(by='class', ascending=True)
        # df = results.pandas().xyxy[0]
        print(df)

        cv.imshow("det_frame", image)
        cv.waitKey(1)

        end = time.time()
        fps = 1 / (end - start)
        # print(f"fps:{fps:.2f}")
